from functools import reduce
import math
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.backend as K
from tensorflow.keras import layers
from losses_utils import *

def YoloLoss(targets, predictions, all_anchors, strides, wh_list, anchor_t=4.0, gr=1.0, class_smooth=0.05):
    lcls = tf.zeros(1, dtype=tf.float32)
    lbox = tf.zeros(1, dtype=tf.float32)
    lobj = tf.zeros(1, dtype=tf.float32)
    tcls, tbox, indices, anchors = build_targets(targets, predictions, all_anchors, strides, wh_list, anchor_t)
    balance = [4.0, 1.0, 0.4]
    hyp = {'box': 0.02, 'cls': 0.21638, 'obj': 0.51728}

    for i, prediction in enumerate(predictions):
        b, a, ij = indices[i] #! image_id, anchor_index, grid index
        tobj = tf.zeros(shape=prediction.shape[:3], dtype=prediction.dtype) #! target obj

        nt = b.shape[0] #! number of targets
        if nt:
            baij = tf.concat([b, a, ij],axis=-1) #! (baij, 3), 其中大概率存在索引
            matched_pred = tf.gather_nd(prediction, baij) #! (baij, no)
            # print(matched_pred.shape)
            pxy, pwh, _, pcls = tf.split(matched_pred, [2, 2, 1, -1], axis=-1)
            # print(pcls.shape)

            #! Regression
            pbox = tf.concat([pxy, pwh], axis=-1) #! 坐标是还原后的, 还原到imgsz大小上
            # print("tbox:", tbox[i][:9])
            # print("pbox:", pbox[:9])
            # print('-=-'*30)
            iou = bbox_iou(pbox, tbox[i], CIoU=True, xyxy=False)
            # print(iou.shape) #! (baij,)
            lbox += tf.reduce_mean(1.0 - iou)
            # print("lbox shape: ", lbox.shape)

            #! Objectness
            iou = tf.maximum(0, iou)
            if gr < 1: iou = (1.0 - gr) + gr * iou
            tobj = tf.scatter_nd(baij, iou, shape=prediction.shape[:3])
            # print(tobj.shape)
            # print(tf.where(tf.greater(tobj, 0)).shape)

            #! Classification
            if pcls.shape[-1] > 1: #! 1
                ind = tf.concat([
                    tf.range(tcls[i].shape[0], dtype=tf.int64)[:, None], 
                    tf.cast(tcls[i], dtype=tf.int64)], axis=-1)
                t = tf.scatter_nd(ind, tf.ones(shape=(tcls[i].shape[0],)), pcls.shape)
                #! Smooth
                t = class_smooth * (1 - t) + (1 - class_smooth) * t
                #! Calc
                # print("t shape: ", t.shape)
                # print("pcls shape: ", pcls.shape)
                lcls += keras.losses.BinaryCrossentropy(from_logits=True)(t, pcls)
            # print("lcls shape: ",lcls.shape)

        # print("tobj shape: ", tobj.shape)
        # print("prediction[..., 4] shape: ", prediction[..., 4].shape)
        obji = keras.losses.BinaryCrossentropy(from_logits=False)(tobj, prediction[..., 4])
        lobj += obji * balance[i]

    lbox *= hyp['box']
    lobj *= hyp['obj']
    lcls *= hyp['cls']

    bs = tobj.shape[0]

    return (lbox + lobj + lcls) * bs, tf.concat([lbox, lobj, lcls], axis=-1)

def build_targets(targets, predictions, all_anchors, strides, wh_list, anchor_t):
    """
        targets: image_id(0), class(1), cx(2), cy(3), w(4), h(5)
    """
    DTYPE = predictions[0].dtype
    targets = tf.cast(targets, dtype=DTYPE)

    na, nl = all_anchors.shape[:2]
    nt = targets.shape[0]
    tcls, tbox, indices, anch = [], [], [], []
    gain = tf.ones(7, dtype=DTYPE)
    ai = tf.tile(tf.range(na, dtype=DTYPE)[:, None], [1, nt])
    targets = tf.concat([tf.tile(targets[None], [na, 1, 1]), ai[:, :, None]], axis=-1)

    g = 0.5
    off = tf.convert_to_tensor([
        [ 0, 0], [ 1, 0], [ 0, 1], [-1, 0], [ 0,-1],
                # [ 1, 1], [ 1,-1], [-1, 1], [-1,-1],
    ], dtype=DTYPE) * g #! (5, 2)

    for i in range(nl):
        # anchors = all_anchors[i] #! (na, 2)
        anchors = all_anchors[i] / strides[i] #! (na, 2) 在feat上的尺寸

        w, h = wh_list[i]
        gain = tf.where(tf.convert_to_tensor([False, False, True, True, True, True, False], dtype=tf.bool),\
            tf.constant([1, 1, w, h, w, h, 1], dtype=gain.dtype), gain)
        
        t = targets * gain #! (3, nt, 7)
        # print(t.shape)
        
        if nt:
            #! Matches
            r = t[:, :, 4:6] / anchors[:, None] #! wh ratio
            j = tf.less(tf.reduce_max(tf.maximum(r, 1/r), axis=-1), anchor_t) #! (3, nt)
            t = tf.gather_nd(t, tf.where(j)) #! (unknown_nt_1, 7)
            # print(t.shape)

            #! Offsets
            gxy = t[:, 2:4] #! cx, cy (unknown_nt_1, 2)
            gxi = gain[2:4] - gxy #! inverse
            j, k = tf.transpose(tf.logical_and(tf.less(gxy % 1, g), tf.greater(gxy, 1)))
            l, m = tf.transpose(tf.logical_and(tf.less(gxi % 1, g), tf.greater(gxi, 1)))
            j = tf.stack([tf.ones_like(j), j, k, l, m])
            t = tf.gather_nd(tf.tile(t[None], [5, 1, 1]), tf.where(j)) #! (unknown_nt_2, 7)
            offsets = tf.gather_nd(tf.zeros_like(gxy)[None] + off[:, None], tf.where(j)) #! (unkown_nt_2, 2)
        else:
            t = targets[0]
            offsets = 0

        #! Define
        b, c, gxy, gwh, a = tf.split(t, [1, 1, 2, 2, 1], axis=-1) #! image_id, class, gxy, gwh, anchors
        gij = tf.cast(gxy - offsets, dtype=tf.int64)
        gi, gj = tf.transpose(gij) #! x坐标, y坐标
        # ij = (tf.clip_by_value(gi, 0, w-1) * w + tf.clip_by_value(gj, 0, h-1))[:, None] #!? 标注一下, 可能出错?
        ij = (tf.clip_by_value(gj, 0, h-1) * h + tf.clip_by_value(gi, 0, w-1))[:, None] #!? 标注一下, 可能出错?
        
        #! Append
        # print(tf.concat([gxy - tf.cast(gij, dtype=gxy.dtype), gwh], axis=-1))
        # print(tf.gather_nd(anchors, tf.cast(a, dtype=tf.int64)))

        indices.append([tf.cast(b, tf.int64), tf.cast(a, tf.int64), ij]) #! 索引
        # tbox.append(tf.concat([gxy - tf.cast(gij, dtype=gxy.dtype), gwh], axis=-1)) #! cx,cy偏移量 wh倍数
        tbox.append(strides[i] * tf.concat([gxy, gwh], axis=-1)) #! cx,cy, w, h在imgsz上的绝对值
        anch.append(tf.gather_nd(anchors, tf.cast(a, dtype=tf.int64)))
        tcls.append(c)

    return tcls, tbox, indices, anch